import os
import re
import subprocess
import sys
from time import sleep

headers = [("User-Agent",
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/75.0.3770.100 Safari/537.36")]

aria2_header_params = " ".join('--header="{}"'.format(x) for x in [": ".join(x) for x in headers])


aria2_path = os.path.join(os.path.dirname(__file__), "aria2c.exe") if sys.platform == "win32" else "aria2c"
aira2_cmd = '%s -x 16 -s 16 -j 64 -k 2M --file-allocation=falloc --disk-cache=128M {header} "{url:}" -d "{dirname:}" -o "{filename:}"' % (
    aria2_path,)

LENGTH = 80

RE_SPEED = re.compile(r'(\d*?\.?\d*?)([KMG])iB/(\d*?\.?\d*?)([KMG])iB\((\d+)%\).*?DL:(\d*?\.?\d*?)([KMG])iB.*?ETA:(.*?)]')
RE_AVESPEED = re.compile(r'\|\s*?([\S]*?)([KM])iB/s\|')


class DownloadFailed(Exception):
    pass


def aria2_download_file(url, filename, dirname='.'):
    cnt = 0
    while cnt < 3:
        try:
            cmd = aira2_cmd.format(url=url, dirname=dirname, filename=filename, header=aria2_header_params)
            p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, universal_newlines=True, encoding='utf8')
            lines = ''
            while p.poll() is None:
                line = p.stdout.readline().strip()
                # if filename.endswith('.mp4') and line:
                if len(line) > 0:
                    lines += line
                    match = RE_SPEED.search(line)
                    if match:
                        got, got_unit, size, size_unit, percent, speed, speed_unit, eta = match.groups()
                        percent = float(percent)
                        speed = float(speed)
                        if speed_unit == 'K':
                            speed /= 1024
                        progress_bar = min(int(LENGTH * percent / 100), LENGTH)
                        print(
                            '\r  |-[' + progress_bar * '*' + (LENGTH - progress_bar) * '.' + '] {:.0f}% {:.2f}M/s'.format(percent, speed),
                            end=' (ctrl+c中断)')
            if p.returncode != 0:
                cnt += 1
                if cnt == 1:
                    clear_files(dirname, filename)
                    sleep(0.16)
            else:
                match = RE_AVESPEED.search(lines)
                if match:
                    ave_speed, unit = match.groups()
                    ave_speed = float(ave_speed)
                    if unit == 'K':
                        ave_speed /= 1024
                    print('\r  |-[' + LENGTH * '*' + '] {:.0f}% {:.2f}M/s'.format(100, ave_speed),
                          end='  (完成)    \n')
                return
        finally:
            p.kill()  # 保证子进程已终止
    clear_files(dirname, filename)
    raise DownloadFailed("download failed")


def clear_files(dirname, filename):
    filepath = os.path.join(dirname, filename)
    if os.path.exists(filepath):
        os.remove(filepath)
    if os.path.exists(filepath + '.aria2'):
        os.remove(filepath + '.aria2')


